import os
import pandas as pd
import pymysql
from dotenv import load_dotenv
from sqlalchemy import create_engine
import random

load_dotenv(override=True)

# 数据库连接配置
db_config = {
    'host': os.getenv("HOST"),
    'port': int(os.getenv("PORT")),
    'user': os.getenv("USER"),
    'password': os.getenv("MYSQL_PW"),
    'database': os.getenv("DB_NAME")
}

print(db_config)

# 获取当前工作目录
current_path = os.getcwd()
print("当前工作目录:", current_path)

# 或者使用pathlib（Python 3.4+）
from pathlib import Path
current_path = Path.cwd()
print("当前工作目录:", current_path)


def replace_user_data():

    # 创建数据库连接
    engine = create_engine(
        f"mysql+pymysql://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['database']}")

    try:
        # 1. 从MySQL表中读取数据
        query = "SELECT * FROM insight_general_table WHERE LABEL_IS_DIFF = 0"
        df_mysql = pd.read_sql(query, engine)

        # 2. 读取CSV文件
        user_product_df = pd.read_csv('/data/user_product.csv')
        group_user_df = pd.read_csv('/data/group_user.csv')

        # 3. 处理用户产品数据（一对多关系）
        # 为每个用户随机选择一个产品
        user_product_mapping = {}
        for user_id, group in user_product_df.groupby('User ID'):
            products = group['Product Name'].tolist()
            user_product_mapping[user_id] = random.choice(products)

        # 4. 处理用户群组数据（多对一关系）
        # 创建用户到群组的映射
        group_mapping = dict(zip(group_user_df['User ID'], group_user_df['Group Name']))

        # 5. 替换数据
        replaced_count = 0

        for index, row in df_mysql.iterrows():
            user_id = row['USER_ID']

            # 如果用户ID在CSV文件中存在，则进行替换
            if user_id in user_product_mapping and user_id in group_mapping:
                # 替换USER_ID（保持不变，但验证存在性）
                # 替换ATTR_PRODUCT_NAME
                df_mysql.at[index, 'ATTR_PRODUCT_NAME'] = user_product_mapping[user_id]
                # 替换ATTR_GROUP_NAME
                df_mysql.at[index, 'ATTR_GROUP_NAME'] = group_mapping[user_id]
                replaced_count += 1

        print(f"成功替换了 {replaced_count} 条记录")

        # 6. 将更新后的数据保存回数据库（可选）
        # 如果需要更新原表，可以使用以下代码：
        # for index, row in df_mysql.iterrows():
        #     update_query = f"""
        #         UPDATE insight_general_table
        #         SET ATTR_PRODUCT_NAME = '{row['ATTR_PRODUCT_NAME']}',
        #             ATTR_GROUP_NAME = '{row['ATTR_GROUP_NAME']}'
        #         WHERE USER_ID = '{row['USER_ID']}' AND LABEL_IS_DIFF = 0
        #     """
        #     # 执行更新操作

        # 或者创建新表保存结果
        df_mysql.to_sql('insight_general_table_updated', engine, if_exists='replace', index=False)
        print("数据已保存到新表: insight_general_table_updated")

        return df_mysql

    except Exception as e:
        print(f"处理过程中出现错误: {e}")
        return None


# 更高效的版本，使用批量操作
def replace_user_data_efficient():

    engine = create_engine(
        f"mysql+pymysql://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['database']}")

    try:
        # 读取数据
        query = "SELECT * FROM insight_general_table WHERE LABEL_IS_DIFF = 0"
        df_mysql = pd.read_sql(query, engine)

        # 读取CSV文件
        user_product_df = pd.read_csv('/data/user_product.csv')
        group_user_df = pd.read_csv('/data/group_user.csv')

        # 创建映射字典
        # 用户-产品映射（随机选择）
        user_product_mapping = user_product_df.groupby('User ID')['Product Name'].apply(
            lambda x: random.choice(x.tolist()) if len(x) > 0 else None
        ).to_dict()

        # 用户-群组映射
        group_mapping = dict(zip(group_user_df['User ID'], group_user_df['Group Name']))

        # 使用向量化操作进行替换
        mask = df_mysql['USER_ID'].isin(user_product_mapping.keys()) & df_mysql['USER_ID'].isin(group_mapping.keys())

        df_mysql.loc[mask, 'ATTR_PRODUCT_NAME'] = df_mysql.loc[mask, 'USER_ID'].map(user_product_mapping)
        df_mysql.loc[mask, 'ATTR_GROUP_NAME'] = df_mysql.loc[mask, 'USER_ID'].map(group_mapping)

        print(f"成功替换了 {mask.sum()} 条记录")

        # 保存结果
        df_mysql.to_sql('insight_general_table_updated', engine, if_exists='replace', index=False)
        print("数据已保存到新表: insight_general_table_updated")

        return df_mysql

    except Exception as e:
        print(f"处理过程中出现错误: {e}")
        return None


# 如果需要直接更新原表
def update_original_table():

    connection = pymysql.connect(**db_config)

    try:
        with connection.cursor() as cursor:
            # 读取CSV文件
            user_product_df = pd.read_csv('data/user_product.csv')
            group_user_df = pd.read_csv('data/group_user.csv')

            # 处理用户产品数据
            user_product_mapping = {}
            for user_id, group in user_product_df.groupby('User ID'):
                products = group['Product Name'].tolist()
                random.seed(hash(user_id))
                user_product_mapping[user_id] = random.choice(products)

            # 处理用户群组数据
            group_mapping = dict(zip(group_user_df['User ID'], group_user_df['Group Name']))

            # 获取所有可用的新用户ID
            available_user_ids = list(set(user_product_df['User ID'].tolist()) & set(group_user_df['User ID'].tolist()))

            # 获取需要更新的所有旧用户ID
            cursor.execute("SELECT DISTINCT USER_ID FROM insight_general_table WHERE LABEL_IS_DIFF = 0")
            old_user_ids = [row[0] for row in cursor.fetchall()]

            # 创建旧用户ID到新用户ID的映射
            old_to_new_user_id = {}
            for i, old_user_id in enumerate(old_user_ids):
                new_user_id = available_user_ids[i % len(available_user_ids)]
                old_to_new_user_id[old_user_id] = new_user_id

            # 批量更新
            update_count = 0
            for old_user_id, new_user_id in old_to_new_user_id.items():
                product_name = user_product_mapping.get(new_user_id, '')
                group_name = group_mapping.get(new_user_id, '')

                update_query = """
                               UPDATE insight_general_table
                               SET USER_ID           = %s,
                                   ATTR_PRODUCT_NAME = %s,
                                   ATTR_GROUP_NAME   = %s
                               WHERE USER_ID = %s \
                                 AND LABEL_IS_DIFF = 0 \
                               """
                cursor.execute(update_query, (new_user_id, product_name, group_name, old_user_id))
                update_count += cursor.rowcount

            connection.commit()
            print(f"成功更新了 {update_count} 条记录")

    except Exception as e:
        print(f"处理过程中出现错误: {e}")
        connection.rollback()
        import traceback
        traceback.print_exc()
    finally:
        connection.close()


if __name__ == "__main__":
    # 使用方法1：创建新表保存结果
    # result_df = replace_user_data_efficient()

    # 使用方法2：直接更新原表（谨慎使用）
    update_original_table()

    # if result_df is not None:
    #     print("前5条处理后的数据:")
    #     print(result_df.head())